import ollama
import re


def get_answer(question: str) -> str:
    """
    使用Ollama官方API调用本地deepseek-r1:8b模型，
    返回去除思考过程及删除``之间内容的回答

    参数:
        question: 提问的文本内容

    返回:
        处理后的模型回答
    """
    try:
        # 调用本地Ollama服务中的模型
        response = ollama.chat(
            model="deepseek-r1:8b",
            messages=[
                {
                    "role": "user",
                    "content": f"请回答以下问题：{question}"
                }
            ]
        )

        # 提取原始回答内容
        raw_answer = response["message"]["content"].strip()

        # 使用正则表达式删除``之间的内容（包括标记本身）
        # 正则表达式解释：
        # \{\{.*?\}\} 匹配`{{`和`}}`之间的任何内容（非贪婪模式）
        # re.DOTALL 使.匹配包括换行符在内的所有字符

        pattern = re.escape("<think>") + ".*?" + re.escape("</think>")
        filtered_answer = re.sub(pattern, '', raw_answer, flags=re.DOTALL)
        # filtered_answer = re.sub(r'\{\{.*?\}\}', '', raw_answer, flags=re.DOTALL)
        #
        # # 去除可能产生的多余空行
        # filtered_answer = re.sub(r'\n+', '\n', filtered_answer).strip()
        # 按行分割文本
        lines = filtered_answer.split('\n')

        # 过滤掉包含"简介："的行
        filtered_lines = [line for line in lines if "简介：" not in line]

        # 将过滤后的行重新组合成文本
        filtered_answer = '\n'.join(filtered_lines)

        return filtered_answer

    except ollama.ResponseError as e:
        return f"API响应错误: {str(e)}"
    except ollama.ConnectionError:
        return "无法连接到Ollama服务，请确保服务已启动"
    except Exception as e:
        return f"发生错误: {str(e)}"


# 示例用法
if __name__ == "__main__":
    test_question = "什么是机器学习？"
    answer = get_answer(test_question)
    print(f"问题: {test_question}")
    print(f"处理后的回答:\n{answer}")
